﻿# main.py
import sys
import pandas as pd
from PyQt5.QtWidgets import (
    QApplication, QWidget, QVBoxLayout, QHBoxLayout,
    QPushButton, QTextEdit, QLineEdit, QDateEdit, QLabel
)
from PyQt5.QtCore import QDate, QTimer
from backtest.engine import AnimatedBacktestEngine
from strategy.ma_crossover import MACrossoverStrategy
from utils.plotly_plot import plot_kline_to_json
from ui.web_view import create_web_view

class MainWindow(QWidget):
    def __init__(self):
        super().__init__()
        self.engine = None
        self.timer = None
        self.running = False
        self.paused = False
        self.speed_interval = 200
        self.init_ui()

    def init_ui(self):
        self.setWindowTitle("量化动画回测系统")
        self.resize(1200, 800)
        layout = QVBoxLayout()

        # 控制区
        ctrl_layout = QHBoxLayout()
        self.symbol_input = QLineEdit("600519")
        self.symbol_input.setFixedWidth(100)
        self.start_date = QDateEdit(QDate(2024, 1, 1))
        self.end_date = QDateEdit(QDate(2024, 12, 31))
        self.run_btn = QPushButton("开始回测")
        self.pause_btn = QPushButton("暂停")
        self.resume_btn = QPushButton("继续")
        self.speed_up_btn = QPushButton("快进 ⏩")
        self.slow_down_btn = QPushButton("慢放 ⏪")

        self.pause_btn.setEnabled(False)
        self.resume_btn.setEnabled(False)

        ctrl_layout.addWidget(QLabel("股票:"))
        ctrl_layout.addWidget(self.symbol_input)
        ctrl_layout.addWidget(QLabel("开始:"))
        ctrl_layout.addWidget(self.start_date)
        ctrl_layout.addWidget(QLabel("结束:"))
        ctrl_layout.addWidget(self.end_date)
        ctrl_layout.addWidget(self.run_btn)
        ctrl_layout.addWidget(self.pause_btn)
        ctrl_layout.addWidget(self.resume_btn)
        ctrl_layout.addWidget(self.speed_up_btn)
        ctrl_layout.addWidget(self.slow_down_btn)
        ctrl_layout.addStretch()

        layout.addLayout(ctrl_layout)

        # 图表区
        self.web_view = create_web_view()
        layout.addWidget(self.web_view, stretch=3)

        # 日志区
        self.log_text = QTextEdit()
        self.log_text.setFixedHeight(120)
        self.log_text.setReadOnly(True)
        layout.addWidget(self.log_text)

        self.setLayout(layout)

        # 绑定事件
        self.run_btn.clicked.connect(self.on_run_clicked)
        self.pause_btn.clicked.connect(self.on_pause_clicked)
        self.resume_btn.clicked.connect(self.on_resume_clicked)
        self.speed_up_btn.clicked.connect(self.on_speed_up_clicked)
        self.slow_down_btn.clicked.connect(self.on_slow_down_clicked)

        self.log("✅ 系统启动，点击【开始回测】运行动画回测")

    def log(self, msg: str):
        self.log_text.append(msg)

    def on_run_clicked(self):
        if self.running:
            self.log("⚠️ 回测已在运行中")
            return

        symbol = self.symbol_input.text().strip()
        if not symbol:
            self.log("❌ 股票代码不能为空")
            return

        start_date = self.start_date.date().toString("yyyyMMdd")
        end_date = self.end_date.date().toString("yyyyMMdd")

        self.engine = AnimatedBacktestEngine()
        df = self.engine.add_data(symbol, start_date, end_date)

        if df.empty:
            self.log("❌ 数据获取失败或为空")
            return

        self.engine.add_strategy(MACrossoverStrategy())
        self.log(f"📊 数据加载完成，共 {len(df)} 根K线")

        self.run_animation()

    def run_animation(self):
        self.running = True
        self.timer = QTimer()
        self.timer.timeout.connect(self.on_timer_tick)
        self.timer.start(self.speed_interval)

        self.log("🚀 动画回测已启动")
        self.run_btn.setEnabled(False)
        self.pause_btn.setEnabled(True)
        self.resume_btn.setEnabled(False)

    def on_timer_tick(self):
        if not self.running or self.paused:
            return

        if self.engine.current_idx >= len(self.engine.data) - 1:
            self.log("✅ 回测完成！")
            self.stop_timer()
            self.run_btn.setEnabled(True)
            self.pause_btn.setEnabled(False)
            self.resume_btn.setEnabled(False)
            return

        if self.engine.run_next():
            current_df = self.engine.get_current_data()
            trades_df = self.engine.get_trades()

            buy_signals = trades_df[trades_df['type'] == 'buy']
            sell_signals = trades_df[trades_df['type'] == 'sell']

            json_data = plot_kline_to_json(current_df, buy_signals, sell_signals)
            self.web_view.page().runJavaScript(f"updateKline(`{json_data}`);")

            stats = self.engine.get_stats()
            self.log(f"📈 第 {len(current_df)} 根 | 资金: {stats.get('value', 0):.2f}")

    def on_pause_clicked(self):
        if self.running and not self.paused:
            self.paused = True
            self.timer.stop()
            self.log("⏸️ 已暂停")
            self.pause_btn.setEnabled(False)
            self.resume_btn.setEnabled(True)

    def on_resume_clicked(self):
        if self.running and self.paused:
            self.paused = False
            self.timer.start(self.speed_interval)
            self.log("▶️ 继续回测")
            self.pause_btn.setEnabled(True)
            self.resume_btn.setEnabled(False)

    def on_speed_up_clicked(self):
        self.speed_interval = max(50, self.speed_interval - 50)
        if self.timer and not self.paused:
            self.timer.stop()
            self.timer.start(self.speed_interval)
        self.log(f"⏩ 快进: {self.speed_interval}ms/根")

    def on_slow_down_clicked(self):
        self.speed_interval += 50
        if self.timer and not self.paused:
            self.timer.stop()
            self.timer.start(self.speed_interval)
        self.log(f"⏪ 慢放: {self.speed_interval}ms/根")

    def stop_timer(self):
        if self.timer:
            self.timer.stop()
            self.timer = None

if __name__ == "__main__":
    app = QApplication(sys.argv)
    window = MainWindow()
    window.show()
    sys.exit(app.exec_())